import json
import os
import random
import uuid

import cv2
import numpy as np

# data source : https://aistudio.baidu.com/aistudio/datasetdetail/103078

base_path = "E:/dataSet/qrcode/qr/"
label_dir = "labels/"
image_dir = "imgs/"

output_path = "E:/dataSet/qrcode_cocodata/"
output_label_dir = "labels/"
output_image_dir = "images/"
output_train_dir = "train/"
output_val_dir = "val/"


def translateFromJson(label_file_name):
    with open(base_path + label_dir + label_file_name, 'r') as label_file:
        json_obj = json.load(label_file)
        image_file_name = str(json_obj['imagePath']).replace("./", '')
        print(image_file_name)
        image_file = cv2.imread(base_path + image_dir + image_file_name)
        image_shape = image_file.shape
        image_width = image_shape[1]
        image_height = image_shape[0]
        label_str = ''
        for shape in json_obj['shapes']:
            x_1 = shape['points'][0][0]
            y_1 = shape['points'][0][1]
            x_2 = shape['points'][1][0]
            y_2 = shape['points'][1][1]
            x_max = max(x_1, x_2)
            x_min = min(x_1, x_2)
            y_max = max(y_1, y_2)
            y_min = min(y_1, y_2)
            label_str = label_str + "0 " + str((x_min + x_max) / 2 / image_width) + ' ' + str(
                (y_min + y_max) / 2 / image_height) + ' ' + str((x_max - x_min) / image_width) + ' ' + str(
                (y_max - y_min) / image_height) + '\n'
        label_str = label_str[0:-1]
        # print(label_str)
        train_blur_cnn_kernel_size = random.randint(3, 9)
        gauss_noisy_noisy_sigma = 33
        save_type = random.randint(1, 6)
        if save_type == 1:
            str(uuid.uuid1())
            # origin data
            save(image_file, str(uuid.uuid1()), save_type, label_str)
        if save_type == 2:
            # gauss blur
            gaussian_blur_extend_image = gauss_blur(image_file, train_blur_cnn_kernel_size)
            save(gaussian_blur_extend_image, str(uuid.uuid1()), save_type, label_str)
        if save_type == 3:
            # gauss noisy
            gauss_noisy_image = gauss_noisy(0, gauss_noisy_noisy_sigma, image_file)
            save(gauss_noisy_image, str(uuid.uuid1()), save_type, label_str)
        if save_type == 4:
            # gauss_noisy + gray
            gray_image = gauss_noisy(0, gauss_noisy_noisy_sigma, image_file)
            gray_image = cv2.cvtColor(gray_image, cv2.COLOR_BGR2GRAY)
            save(gray_image, str(uuid.uuid1()), save_type, label_str)
        if save_type == 6:
            # gray
            gray_image = cv2.cvtColor(image_file, cv2.COLOR_BGR2GRAY)
            save(gray_image, str(uuid.uuid1()), save_type, label_str)

        if save_type % 2 == 0:
            # gauss_blur + gray
            train_image = gauss_blur(image_file, train_blur_cnn_kernel_size)
            # noisy_sigma = random.randint(5, 10) * (((image_height // 200) + (image_width // 200)) // 2)
            # train_image = gauss_noisy(0, noisy_sigma, train_image)
            train_image = cv2.cvtColor(train_image, cv2.COLOR_BGR2GRAY)
            save(train_image, str(uuid.uuid1()), save_type, label_str)

        if save_type % 3 == 0:
            # gauss_blur + gauss_noisy + gray
            train_image = gauss_blur(image_file, train_blur_cnn_kernel_size)
            train_image = gauss_noisy(0, gauss_noisy_noisy_sigma, train_image)
            train_image = cv2.cvtColor(train_image, cv2.COLOR_BGR2GRAY)
            save(train_image, str(uuid.uuid1()), save_type, label_str)

        if save_type == 5:
            # detect data
            blur_cnn_kernel_size = random.randint(3, 13)
            save_detect_type = random.randint(1, 6)
            if save_detect_type == 1:
                save(image_file, str(uuid.uuid1()), save_type, label_str)
            elif save_detect_type == 2:
                gauss_blur_image = gauss_blur(image_file, blur_cnn_kernel_size)
                save(gauss_blur_image, str(uuid.uuid1()), save_type, label_str)
            elif save_detect_type == 3:
                gauss_noisy_image = gauss_noisy(0, gauss_noisy_noisy_sigma + 10, image_file)
                save(gauss_noisy_image, str(uuid.uuid1()), save_type, label_str)
            elif save_detect_type == 4:
                gray_image = gauss_noisy(0, gauss_noisy_noisy_sigma + 10, image_file)
                gray_image = cv2.cvtColor(image_file, cv2.COLOR_BGR2GRAY)
                save(gray_image, str(uuid.uuid1()), save_type, label_str)
            elif save_detect_type == 6:
                # gray
                gray_image = cv2.cvtColor(image_file, cv2.COLOR_BGR2GRAY)
                save(gray_image, str(uuid.uuid1()), save_type, label_str)
            elif save_detect_type == 5:
                detect_image = gauss_blur(image_file, blur_cnn_kernel_size)
                # detect_image = gauss_noisy(0, 40, detect_image)
                detect_image = cv2.cvtColor(detect_image, cv2.COLOR_BGR2GRAY)
                save(detect_image, str(uuid.uuid1()), save_type, label_str)


def save(image_file, file_id, save_type, label_str):
    if save_type != 5:
        cv2.imwrite(output_path + output_train_dir + output_image_dir + file_id + '.jpg', image_file)
        with open(output_path + output_train_dir + output_label_dir + file_id + '.txt', 'w') as output_file:
            output_file.write(label_str)
    else:
        cv2.imwrite(output_path + output_val_dir + output_image_dir + file_id + '.jpg', image_file)
        with open(output_path + output_val_dir + output_label_dir + file_id + '.txt', 'w') as output_file:
            output_file.write(label_str)


def gauss_blur(image_file, blur_cnn_kernel_size):
    if blur_cnn_kernel_size % 2 == 0:
        blur_cnn_kernel_size = blur_cnn_kernel_size + 1
    gauss_blur_image = cv2.GaussianBlur(image_file, (blur_cnn_kernel_size, blur_cnn_kernel_size), 0)
    return gauss_blur_image


def gauss_noisy(mean, sigma, image):
    image_shape = image.shape
    image_width = image_shape[1]
    image_height = image_shape[0]
    image_channel = 3
    gauss = np.random.normal(mean, sigma, (image_height, image_width, image_channel))
    noisy_img = image + gauss
    noisy_img = np.clip(noisy_img, a_min=0, a_max=255).astype(np.uint8)
    return noisy_img


if __name__ == '__main__':
    label_file_list = os.listdir(base_path + label_dir)
    index = 0
    total = len(label_file_list)
    for label_file_name in label_file_list:
        translateFromJson(label_file_name)
        index = index + 1
        print("processing:" + str(index) + "/" + str(total))
